import argparse
import yaml

def sed_arg_parser():
    parser = argparse.ArgumentParser("SED system with in-domain and out-of-domain inputs.")
    parser.add_argument(
        "--data_path",
        default="./configs/data_paths.yaml",
        help="The yaml file contains path of .wav files.",
    )
    parser.add_argument(
        "--feat_conf",
        default="./configs/feats.yaml",
        help="The yaml file contains feature configurations.",
    )
    parser.add_argument(
        "--opt_conf",
        default="./configs/opt.yaml",
        help="The yaml file contains training/optimization configurations.",
    )
    parser.add_argument(
        "--log_dir",
        default="./exp/",
        help="Directory where to save tensorboard logs, saved models, etc.",
    )
    parser.add_argument(
        "--gpus",
        default="0,",
        help="In format: 'num' + ',', the comma is required. If you want to use multiple gpus, use comma plus numbers, e.g, '0,1'.",
    )
    parser.add_argument(
        "--resume_from_checkpoint",
        default=None,
        help="Allow the training to be resumed, take as input a previously saved model (.ckpt).",
    )
    parser.add_argument(
        "--test_from_checkpoint", default=None, help="Use model in inference mode only"
    )
    parser.add_argument(
        "--exp_name",
        default="",
        required=True
    )
    return parser

def reconfig(args, unknown_args):
    def arg_type_rec(arg: str):
        arg = arg.strip()
        # list
        if "[" in arg and "]" in arg:
            return eval(arg)
        elif arg in ["True", "False"]:
            return eval(arg)
        elif arg.isnumeric():
            return int(arg)
        # float
        elif arg.replace(".", "", 1).isnumeric():
            return float(arg)
        # int and str
        else:
            return arg
    
    # load yaml configs
    with open(args.data_path, "r") as f:
        data_paths = yaml.safe_load(f)

    with open(args.feat_conf, "r") as f:
        feat_configs = yaml.safe_load(f)

    with open(args.opt_conf, "r") as f:
        opt_configs = yaml.safe_load(f)

    # process unknown args
    unknown_dicts = {}
    for i in range(0, len(unknown_args), 2):
        unknown_dicts[unknown_args[i][2:]] = arg_type_rec(unknown_args[i+1])
    
    for k, v in unknown_dicts.items():
        for config in [data_paths, feat_configs, opt_configs]:
            for k1 in config.keys():
                if type(config[k1]) is dict:
                    if k in list(config[k1].keys()):
                        config[k1][k] = v
                        print(f"Overwrite {k1}-{k} to {v}")
                        break
            if k in list(config.keys()):
                config[k] = v
                print(f"Overwrite {k} to {v}")
                break
        else:
            print(f"Unknown argument {k}={v}")

    # Add root path to all data paths
    for k in data_paths.keys():
        if ("root" not in k) and ("od" not in k):
            data_paths[k] = data_paths["id_root"] + data_paths[k]

    return args, data_paths, feat_configs, opt_configs